package io.sunny.platform.ai.deepseek.controller;

import io.github.pigmesh.ai.deepseek.core.DeepSeekClient;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionRequest;
import io.micrometer.common.util.StringUtils;
import io.sunny.platform.ai.deepseek.config.DeepSeekConfigProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.concurrent.CompletableFuture;

/**
 * @Author SUNNY
 * @Date 2025/11/30 15:51
 * @PackageName: io.sunny.platform.ai.deepseek.controller
 * @Description: 流式响应处理
 * @Version 1.0
 */
@Slf4j
@RestController
@RequestMapping("/api/ai")
public class StreamAIController {
    private final DeepSeekClient deepSeekClient;

    public StreamAIController(DeepSeekConfigProperties config) {
        this.deepSeekClient = new DeepSeekClient.Builder()
                .openAiApiKey(config.getKey())
                .baseUrl(config.getBaseUrl())
                .connectTimeout(config.getTimeout().getConnect())
                .readTimeout(config.getTimeout().getRead())
                .build();
    }

    @GetMapping("/stream-chat")
    public SseEmitter streamChat(@RequestParam String message, @RequestParam(required = false) String sessionId) {
        SseEmitter emitter = new SseEmitter(30000L);   // 30秒超时

        // 设置事件处理器
        emitter.onCompletion(() -> log.debug("SSE连接完成"));
        emitter.onTimeout(() -> log.debug("SSE连接超时"));
        emitter.onError(error -> log.error("SSE连接错误", error));

        // 异步处理流式响应
        CompletableFuture.runAsync(() -> {
            try {
                ChatCompletionRequest request = ChatCompletionRequest.builder()
                        .stream(true)
                        .addUserMessage(message)
                        .user(sessionId)
                        .build();

                deepSeekClient.chatFluxCompletion(request)
                        .doOnSubscribe(subscription -> {
                            sendSseEvent(emitter, "start", "开始生成");
                        })
                        .subscribe(
                                chunk -> {
                                    String content = chunk.choices().get(0).message().content();
                                    if (StringUtils.isNotBlank(content)) {
                                        sendSseEvent(emitter, "data", content);
                                    }
                                    sendSseEvent(emitter, "start", "开始生成");
                                },
                                error -> {
                                    log.error("流式处理异常", error);
                                    sendSseEvent(emitter, "error", "生成过程中发生错误");
                                    emitter.complete();
                                },
                                () -> {
                                    sendSseEvent(emitter, "end", "生成完成");
                                    emitter.complete();
                                }
                        );
            } catch (Exception e) {
                log.error("流式处理异常", e);
                sendSseEvent(emitter, "error", "系统异常");
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }

    private void sendSseEvent(SseEmitter emitter, String event, Object data) {
        try {
            emitter.send(SseEmitter.event()
                    .name(event)
                    .data(data));
        } catch (Exception e) {
            log.warn("SSE事件发送失败", e);
        }
    }
}
